88c502
@@ -22,7 +22,6 @@
 package org.wildfly.clustering.web.undertow.session;
 
 import io.undertow.security.api.AuthenticatedSessionManager.AuthenticatedSession;
-import io.undertow.security.idm.Account;
 import io.undertow.server.HttpServerExchange;
 import io.undertow.server.session.SessionConfig;
 import io.undertow.server.session.SessionListener.SessionDestroyedReason;
@@ -31,6 +30,8 @@
import io.undertow.servlet.handlers.security.CachedAuthenticatedSessionHandler;
 import java.io.NotSerializableException;
 import java.io.Serializable;
 import java.util.AbstractMap.SimpleImmutableEntry;
+import java.util.Arrays;
+import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
@@ -49,6 +50,8 @@
import org.wildfly.clustering.web.session.SessionManager;
 public class DistributableSession implements io.undertow.server.session.Session {
     // Undertow stores the authenticated session in the HttpSession using a special attribute with the following name
     private static final String AUTHENTICATED_SESSION_ATTRIBUTE_NAME = CachedAuthenticatedSessionHandler.class.getName() + ".AuthenticatedSession";
+    // These mechanisms can auto-reauthenticate and thus use local context (instead of replicating)
+    private static final Set<String> AUTO_REAUTHENTICATING_MECHANISMS = new HashSet<>(Arrays.asList(HttpServletRequest.BASIC_AUTH, HttpServletRequest.DIGEST_AUTH, HttpServletRequest.CLIENT_CERT_AUTH));
 
     private final UndertowSessionManager manager;
     private final Batch batch;
@@ -122,8 +125,8 @@
public class DistributableSession implements io.undertow.server.session.Session
         Session<LocalSessionContext> session = this.entry.getKey();
         try (BatchContext context = this.manager.getSessionManager().getBatcher().resumeBatch(this.batch)) {
             if (AUTHENTICATED_SESSION_ATTRIBUTE_NAME.equals(name)) {
-                Account account = (Account) session.getAttributes().getAttribute(name);
-                return (account != null) ? new AuthenticatedSession(account, HttpServletRequest.FORM_AUTH) : session.getLocalContext().getAuthenticatedSession();
+                AuthenticatedSession auth = (AuthenticatedSession) session.getAttributes().getAttribute(name);
+                return (auth != null) ? auth : session.getLocalContext().getAuthenticatedSession();
             }
             return session.getAttributes().getAttribute(name);
         }
@@ -137,17 +140,8 @@
public class DistributableSession implements io.undertow.server.session.Session
         Session<LocalSessionContext> session = this.entry.getKey();
         try (BatchContext context = this.manager.getSessionManager().getBatcher().resumeBatch(this.batch)) {
             if (AUTHENTICATED_SESSION_ATTRIBUTE_NAME.equals(name)) {
-                AuthenticatedSession authSession = (AuthenticatedSession) value;
-                // If using FORM authentication, we store the corresponding Account in a session attribute
-                if (authSession.getMechanism().equals(HttpServletRequest.FORM_AUTH)) {
-                    Account account = (Account) session.getAttributes().setAttribute(name, authSession.getAccount());
-                    return (account != null) ? new AuthenticatedSession(account, HttpServletRequest.FORM_AUTH) : null;
-                }
-                // Otherwise we store the whole AuthenticatedSession in the local context
-                LocalSessionContext localContext = session.getLocalContext();
-                AuthenticatedSession old = localContext.getAuthenticatedSession();
-                localContext.setAuthenticatedSession(authSession);
-                return old;
+                AuthenticatedSession auth = (AuthenticatedSession) value;
+                return AUTO_REAUTHENTICATING_MECHANISMS.contains(auth.getMechanism()) ? this.setLocalContext(auth) : session.getAttributes().setAttribute(name, auth);
             }
             if (!(value instanceof Serializable)) {
                 throw new IllegalArgumentException(new NotSerializableException(value.getClass().getName()));
@@ -167,14 +161,8 @@
public class DistributableSession implements io.undertow.server.session.Session
         Session<LocalSessionContext> session = this.entry.getKey();
         try (BatchContext context = this.manager.getSessionManager().getBatcher().resumeBatch(this.batch)) {
             if (AUTHENTICATED_SESSION_ATTRIBUTE_NAME.equals(name)) {
-                Account account = (Account) session.getAttributes().removeAttribute(name);
-                if (account != null) {
-                    return new AuthenticatedSession(account, HttpServletRequest.FORM_AUTH);
-                }
-                LocalSessionContext localContext = session.getLocalContext();
-                AuthenticatedSession old = localContext.getAuthenticatedSession();
-                localContext.setAuthenticatedSession(null);
-                return old;
+                AuthenticatedSession auth = (AuthenticatedSession) session.getAttributes().removeAttribute(name);
+                return (auth != null) ? auth : this.setLocalContext(null);
             }
             Object old = session.getAttributes().removeAttribute(name);
             if (old != null) {
@@ -218,4 +206,11 @@
public class DistributableSession implements io.undertow.server.session.Session
             return id;
         }
     }
+
+    private AuthenticatedSession setLocalContext(AuthenticatedSession auth) {
+        LocalSessionContext localContext = this.entry.getKey().getLocalContext();
+        AuthenticatedSession old = localContext.getAuthenticatedSession();
+        localContext.setAuthenticatedSession(auth);
+        return old;
+    }
 }
